import torch

from preprocess.traindataset import TrainDataset
from torch.utils.data import DataLoader
from model.Unet3d import Unet3d
from utils.loss import DiceLoss,TverskyLoss

from tqdm import trange
from tensorboardX import SummaryWriter
writer = SummaryWriter()

epoch = 100

# 初始化数据集
train_dataset = TrainDataset()

# 定义数据加载
train_loader = DataLoader(dataset=train_dataset, batch_size=8, shuffle=True)

# 初始化模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Unet3d(in_channel=1, out_channel=1, training=True).to(device)

# 初始化损失函数
# loss_fuc = DiceLoss()
loss_fuc = TverskyLoss()

# 初始化优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

model.train()
for n_iter in trange (epoch):
    total_loss = 0
    num = 0
    for step, (image, mask) in enumerate(train_loader):
        # print(image.shape,mask.shape)
        image = image .to(device)
        mask = mask.to(device)
        mask_pre = model(image)
        # print (mask.shape ,mask_pre.shape)
        loss = loss_fuc(mask.unsqueeze(1), mask_pre)
        loss.backward()
        optimizer.step()
        total_loss +=loss.item()
        num+=1
    # writer.add_image('Image', image[0,0,16,:,:].unsqueeze(0), n_iter)
    writer.add_image('Pred', mask_pre[0,0,16,:,:].unsqueeze(0), n_iter)
    writer.add_image('Mask', mask[0,16,:,:].unsqueeze(0), n_iter)
    print(total_loss/num)

